/*
 * This is written by Lulu Han and Yunguo Guan.
 * E-mail: locomotive_crypto@163.com
 */ 

#ifndef PPPOI_SHE_H
#define PPPOI_SHE_H

#include  <NTL/ZZ.h>

class SHE{
public:
	
	class PrivateKey{
	public:
		PrivateKey(){}
		PrivateKey(const NTL::ZZ& p, const NTL::ZZ& q, const NTL::ZZ& L);
		
		NTL::ZZ get_N() const;
		NTL::ZZ get_p() const;
		NTL::ZZ get_q() const;
		NTL::ZZ get_L() const;
	
	private:
		NTL::ZZ m_N;
		NTL::ZZ m_p;
		NTL::ZZ m_q;
		NTL::ZZ m_L;
	};
	
	
	class PublicKey{
	public:
		PublicKey(){}
		/*
		 * Note that e0 is the encryption of plaintext 0, 
		 * 		e1 is the encryption of plaintext 1,
		 * 		en1 is the encryption of -1.
		 */
		PublicKey(const NTL::ZZ& N, const NTL::ZZ& e01, const NTL::ZZ& e02, const NTL::ZZ& e1, 
			const NTL::ZZ& en1, long k0, long k1, long k2);
		
		NTL::ZZ get_N() const;
		long get_k2() const;
		long get_k1() const;
		long get_mul_depth() const;
		
		NTL::ZZ get_enc01() const;
		NTL::ZZ get_enc02() const;
		NTL::ZZ get_enc1() const;
		NTL::ZZ get_encn1() const;
	private:
		NTL::ZZ m_N;
		NTL::ZZ m_enc_01;
		NTL::ZZ m_enc_02;
		NTL::ZZ m_enc_1;
		NTL::ZZ m_enc_neg_1;
		long m_k0, m_k1, m_k2;
	};
	
	class Encryptor {
	public:
		Encryptor(const PublicKey& pk);
		
		void encrypt(const NTL::ZZ& pt, NTL::ZZ& ct);
		
		void he_add_cc(const NTL::ZZ& ct1, const NTL::ZZ& ct2, NTL::ZZ& ct);
		void he_add_cp(const NTL::ZZ& ct1, const NTL::ZZ& pt1, NTL::ZZ& ct);
		
		void he_mul_cc(const NTL::ZZ& ct1, const NTL::ZZ& ct2, NTL::ZZ& ct);
		void he_mul_cp(const NTL::ZZ& ct1, const NTL::ZZ& pt1, NTL::ZZ& ct);
	private:
		PublicKey m_pk;
	};
	
	class Decryptor{
	public:
		Decryptor(const PrivateKey& sk);
		
		void decrypt(const NTL::ZZ& ct, NTL::ZZ& pt);
	
	private:
		PrivateKey m_sk;
	};
	
	// Generate a random key pair
	static void key_gen(PrivateKey& sk, PublicKey& pk, long k0, long k1, long k2);
};


#endif // PPPOI_SHE_H
